{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing BABEL labels\n",
    "[BABEL](https://babel.is.tue.mpg.de/) labels mocap sequences from [AMASS](https://amass.is.tue.mpg.de) with action labels. \n",
    "A single sequence in BABEL can have multiple action labels associated with it, from multiple annotators. \n",
    "Here, we present code to load data from BABEL, visualize the mocap sequence rendered as a 2D video, and view the action labels corresponding to the sequence. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preparing the environment\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib notebook\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from os.path import join as ospj\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import pprint\n",
    "pp = pprint.PrettyPrinter()\n",
    "\n",
    "from IPython.display import HTML"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load BABEL\n",
    "We assume that you have downloaded BABEL annotations from the [website](https://babel.is.tue.mpg.de/data.html) and placed the downloaded `babel_v1.0_release` folder in `data/`. The BABEL data is provided as two sets -- BABEL dense and BABEL extra. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "d_folder = '../data/babel_v1.0_release'  # Data folder\n",
    "l_babel_dense_files = ['train', 'val', 'test']\n",
    "l_babel_extra_files = ['extra_train', 'extra_val']\n",
    "\n",
    "# BABEL Dataset \n",
    "babel = {}\n",
    "for file in l_babel_dense_files:\n",
    "    babel[file] = json.load(open(ospj(d_folder, file+'.json')))\n",
    "    \n",
    "for file in l_babel_extra_files:\n",
    "    babel[file] = json.load(open(ospj(d_folder, file+'.json')))    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### View random annotation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let us view an annotation data structure from the BABEL. \n",
    "The overall data structure is a dictionary, with a unique sequence ID as key and the annotation as value.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_random_babel_ann():\n",
    "    '''Get annotation from random sequence from a random file'''\n",
    "    file = np.random.choice(l_babel_dense_files + l_babel_extra_files)\n",
    "    seq_id = np.random.choice(list(babel[file].keys()))\n",
    "    print('We are visualizing annotations for seq ID: {0} in \"{1}.json\"'.format(seq_id, file))\n",
    "    ann = babel[file][seq_id]\n",
    "    return ann, file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "We are visualizing annotations for seq ID: 3312 in \"test.json\"\n",
      "{'babel_sid': 3312,\n",
      " 'dur': 76.73,\n",
      " 'feat_p': 'CMU/CMU/86/86_08_poses.npz',\n",
      " 'frame_ann': {'anntr_id': 'c6065e9c-1652-46df-a45f-fe8b8158428f',\n",
      "               'babel_lid': 'a642048f-7fa9-402f-a4c1-d7e9e7f696d1',\n",
      "               'labels': [{'act_cat': None,\n",
      "                           'end_t': 68.093,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': 'ad703788-bd17-42d4-854b-2b64cb58ee16',\n",
      "                           'start_t': 59.51},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 32.82,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': '1785aeca-53ce-4a33-a249-8a5d3466ea95',\n",
      "                           'start_t': 27.445},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 52.426,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': '12768b82-b342-46ee-ae60-e158f8b1dd47',\n",
      "                           'start_t': 47.843},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 59.51,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': '435bd5a6-01e9-4fc4-abee-642954466832',\n",
      "                           'start_t': 53.26},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 40.007,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': 'd3911406-ad83-4438-941c-919bf296d5e1',\n",
      "                           'start_t': 33.382},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 76.733,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': 'f222a4d9-a8d5-4002-893b-4df102e1e0fa',\n",
      "                           'start_t': 70.593},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 2.252,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': '35e605ec-c9f8-4c9d-8320-680de71837ce',\n",
      "                           'start_t': 0.294},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 6.961,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': 'fdaead4c-0a37-4579-a42a-4a94145570b9',\n",
      "                           'start_t': 4.232},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 70.593,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': '52d3c3e9-102b-4cf0-b082-cd416a7b5f64',\n",
      "                           'start_t': 68.093},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 4.232,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': 'f524e2df-36e2-45ce-a54e-892fdb7353d0',\n",
      "                           'start_t': 2.252},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 9.336,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': '7f265bed-f445-4b6b-a41f-c62106d7be3b',\n",
      "                           'start_t': 6.961},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 47.843,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': '1aa33355-a669-45a6-86a9-19ae862a47e9',\n",
      "                           'start_t': 40.007},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 15.523,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': 'd9c310f5-fc1e-47d8-b2f7-075c31a2eb6d',\n",
      "                           'start_t': 9.523},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 22.507,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': 'f7a71a16-2807-49f7-8a66-7df3e678e161',\n",
      "                           'start_t': 15.523},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 0.294,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': '3f57a657-2c8f-4995-87a4-965bcf8ea2a6',\n",
      "                           'start_t': 0},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 9.523,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': 'c9f97199-97eb-463c-a04e-a511413ad5ba',\n",
      "                           'start_t': 9.336},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 33.382,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': 'dac4fabe-e96c-411c-ad2e-29211e8c212a',\n",
      "                           'start_t': 32.82},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 53.26,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': 'ed99bf22-3ea5-45a6-9df3-17e67e49f119',\n",
      "                           'start_t': 52.426},\n",
      "                          {'act_cat': None,\n",
      "                           'end_t': 27.445,\n",
      "                           'proc_label': None,\n",
      "                           'raw_label': None,\n",
      "                           'seg_id': '5c459b13-35e6-4c36-8ec4-9eb1536bfe95',\n",
      "                           'start_t': 22.507}],\n",
      "               'mul_act': True},\n",
      " 'seq_ann': {'anntr_id': 'a217bb6b-93ae-4611-8e53-d4318ed5be00',\n",
      "             'babel_lid': '037dc092-28d5-4537-9632-9a91fc9f7fb9',\n",
      "             'labels': [{'act_cat': None,\n",
      "                         'proc_label': None,\n",
      "                         'raw_label': None,\n",
      "                         'seg_id': 'f7d4b8fa-de77-487f-a08c-84bbc05c3148'}],\n",
      "             'mul_act': True},\n",
      " 'url': 'https://babel-renders.s3.eu-central-1.amazonaws.com/003312.mp4'}\n"
     ]
    }
   ],
   "source": [
    "ann, _ = get_random_babel_ann()\n",
    "pp.pprint(ann)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the action labels from `test.json` are not available publicly. \n",
    "Also note that the internal data structures of BABEL dense and BABEL extra differ slightly. \n",
    "For a detailed description of the annotation, see [BABEL's data page](https://babel.is.tue.mpg.de/data.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize a mocap seq. and its action labels "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_vid_html(url):\n",
    "    '''Helper code to embed a URL in a notebook'''\n",
    "    html_code = '<div align=\"middle\"><video width=\"80%\" controls>'\n",
    "    html_code += f'<source src=\"{url}\" type=\"video/mp4\">'\n",
    "    html_code += '</video></div>'\n",
    "    return html_code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_labels(ann, file):\n",
    "    # Get sequence labels and frame labels if they exist\n",
    "    seq_l, frame_l = None, None\n",
    "    if 'extra' not in file:\n",
    "        if ann['seq_ann'] is not None:\n",
    "            seq_l = [seg['raw_label'] for seg in ann['seq_ann']['labels']]\n",
    "        if ann['frame_ann'] is not None:\n",
    "            frame_l = [(seg['raw_label'], seg['start_t'], seg['end_t']) for seg in ann['frame_ann']['labels']]\n",
    "    else:\n",
    "        # Load labels from 1st annotator (random) if there are multiple annotators\n",
    "        if ann['seq_anns'] is not None:\n",
    "            seq_l = [seg['raw_label'] for seg in ann['seq_anns'][0]['labels']]\n",
    "        if ann['frame_anns'] is not None:\n",
    "            frame_l = [(seg['raw_label'], seg['start_t'], seg['end_t']) for seg in ann['frame_anns'][0]['labels']]\n",
    "    return seq_l, frame_l"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Visualize a random mocap and its annotation from BABEL, by running the cell below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "We are visualizing annotations for seq ID: 7536 in \"train.json\"\n",
      "Sequence labels:  ['pace and shake hand']\n",
      "Frame labels: (action label, start time, end time)\n",
      "[('walk', 0, 2.106),\n",
      " ('transition', 2.106, 2.845),\n",
      " ('make a knocking gesture', 2.845, 3.507),\n",
      " ('transition', 3.466, 4.6),\n",
      " ('turn around', 4.519, 5.519),\n",
      " ('walk back', 5.424, 7.734)]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div align=\"middle\"><video width=\"80%\" controls><source src=\"https://babel-renders.s3.eu-central-1.amazonaws.com/007536.mp4\" type=\"video/mp4\"></video></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ann, file = get_random_babel_ann()\n",
    "seq_l, frame_l = get_labels(ann, file)\n",
    "print('Sequence labels: ', seq_l)\n",
    "print('Frame labels: (action label, start time, end time)')\n",
    "pp.pprint(frame_l)       \n",
    "HTML(get_vid_html(ann['url']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- If you are interested in loading the mocap sequence in 3D, please refer to the tutorials in [AMASS](https://github.com/nghorbani/amass/tree/master/notebooks)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
